# -*- coding: utf-8 -*-
"""Centrality.py
flow through centrality was developed in
K.S. Gomez-Haibach and M.A. Gomez: Revised Centrality Measures Tell a Robust Story of Ion Conduction in Solids. 
The Journal of Physical Chemistry B 127, 9258 (2023).
However, it uses terms in our original centrality paper
R. Krueger, F. Haibach, D. Fry and M. Gomez: Centrality measures highlight proton traps and access points to proton highways in kinetic Monte Carlo trajectories. Journal of Chemical Physics 142, 154110 (2015).
which is refferred to in the comments
Code from Gomez group at Mount Holyoke College - adapted for this tutorial 2025
"""

import os
import numpy

#setting up arrays, reading Energies.binding file
numberOfBindingSites=96 #changes based on system
posX=numpy.zeros(numberOfBindingSites)
posY=numpy.zeros(numberOfBindingSites)
posZ=numpy.zeros(numberOfBindingSites)
Energies=numpy.zeros(numberOfBindingSites)
protonNumber=numpy.zeros(numberOfBindingSites,dtype=int)

isite=0
#temperature
temperature=1000.#Kelvin
#1.3806503e-23 is Boltzmann's constant in units of J/K
#6.24150974e18 is the conversion from J to eV
kb = 1.3806503e-23*6.24150974e18;
beta = 1/(kb*temperature)  # 1/eV units

sumNorm=0.0
min=0.0
nameoffile="Energies.binding"
for line in open(nameoffile,'r'):
	dummy=line.split()
	protonNumber[isite]=int(eval(dummy[0]))-1
	posX[isite]=eval(dummy[1])
	posY[isite]=eval(dummy[2])
	posZ[isite]=eval(dummy[3])
	Energies[isite]=eval(dummy[4])
	if (min>Energies[isite]):
		min=Energies[isite]
		imin=isite
	isite=isite+1

print("Min found at site")
print(imin+1)

print("Number of Sites found")
print(isite)

#using lowest site energy to shift all energies relative to it
#calculate Boltzmann probabilities
for i in range(numberOfBindingSites):
	#print(i+1,Energies[i])
	Energies[i]=Energies[i]-min
	sumNorm=sumNorm+numpy.exp(-beta*Energies[i]) #normalization factor

sitesFile=open("Sites",'w')
#IIA2a in paper
Boltz=numpy.zeros(numberOfBindingSites)
#print("Boltzmann Distribution")
for i in range(numberOfBindingSites):
	if (sumNorm>0):
	  Boltz[i]=numpy.exp(-beta*Energies[i])/sumNorm
	else:
	  Boltz[i]=0.
	  #When frequencies are available, would need to divide by frequency
	  #product at the lth site.  See Centrality Paper Section II A .2.(a)
	  #print(i+1,Boltz[i],Energies[i])
	  sitesFile.write(str(i+1)+" "+str(posX[i])+" "+str(posY[i])+" "+str(posZ[i])+" "+str(Boltz[i])+"\n")
sitesFile.close()


# W - matrix whose rows are the boltmann distribution (just below eqn 1)
W=numpy.zeros((numberOfBindingSites,numberOfBindingSites))
for i in range(numberOfBindingSites):
	for j in range(numberOfBindingSites):
		W[i][j]=Boltz[j]

#print("W")
#print(W)

#Reading transition state information
ETS=numpy.zeros((numberOfBindingSites,numberOfBindingSites))
RateConstant=numpy.zeros((numberOfBindingSites,numberOfBindingSites))
nameoffile="Energies.TS"
for line in open(nameoffile,'r'):
	dummy=line.split()
	i=int(eval(dummy[0]))-1
	j=int(eval(dummy[1]))-1
	ETS[i][j]=ETS[j][i]=eval(dummy[3])-min
	#print(i+1,j+1,ETS[i][j]+min)
	#IIA1c in paper
	RateConstant[i][j]=numpy.exp(-beta*(ETS[i][j]-Energies[i]))
	RateConstant[j][i]=numpy.exp(-beta*(ETS[j][i]-Energies[j]))
	#If frequencies are available, rate constants wold be multipied by the 
	#product of frequencies at the minimum divided by the product of the 
	#frequencies at the transition state as in the paper section II.A .1 (c)
	#print(i+1,j+1,RateConstant[i][j],RateConstant[j][i])

print("RateConstant")
for i in range(numberOfBindingSites):
	for j in range(numberOfBindingSites):
		if (RateConstant[i][j]>0):
			 print(i+1,j+1,RateConstant[i][j])
		if (RateConstant[j][i]>0):
			 print(j+1,i+1,RateConstant[j][i])


connectionsFile=open("Connections",'w')
print("Probabilities")
#IIA2c
probability=numpy.zeros((numberOfBindingSites,numberOfBindingSites))
for i in range(numberOfBindingSites):
	sum=0.0
	for j in range(numberOfBindingSites):
		sum=sum+RateConstant[i][j]
	if sum>0.0:
	  for j in range(numberOfBindingSites):
                  probability[i][j]=RateConstant[i][j]/sum
                  if probability[i][j]>0:
                          print(i+1,j+1,probability[i][j])
                          connectionsFile.write(str(i+1)+" "+str(j+1)+" "+str(probability[i][j])+" "+str(RateConstant[i][j])+"\n")
connectionsFile.close()

#IIA2c - c_n in eq. 1 is the expected time for first step starting at n 
cVector=numpy.zeros(numberOfBindingSites)
#Expected time for first step average over all in is the last sum  in equation 1 
expectedTimeForFirstStep=0.0
for n in range(numberOfBindingSites):
	sum=0.0
	for m in range(numberOfBindingSites):
		if probability[n][m]>0:
			sum=sum+probability[n][m]/RateConstant[n][m]
	cVector[n]=sum
	expectedTimeForFirstStep=expectedTimeForFirstStep+Boltz[n]*cVector[n]

print("cVector")
print(cVector)
print("expectedTimeForFirstStep")
print(expectedTimeForFirstStep)
#Fundamental Matrix Inverse I-P+W
FundamentalMatrixInverse=numpy.zeros((numberOfBindingSites,numberOfBindingSites))
FundamentalMatrixInverse=numpy.identity(numberOfBindingSites)-probability+W

print("probability")
print(probability)
print("W")
print(W)
print("Fundamental Matrix Inverse")
print(FundamentalMatrixInverse)

#If an oxygen is removed from the system, all the binding sites on it are removed from the graph
#and the related matrices and vectors - in the commented out lines that follow, you can see how to implement
#Now that all is together with right correspondances, we are going to remove rows and columns for 
#casesToRemove=numpy.zeros(4,dtype=int)
#casesToRemove[0]=35
#casesToRemove[1]=15
#casesToRemove[2]=17
#casesToRemove[3]=49
#casesToRemove=casesToRemove-1

#removing cases
#FundamentalMatrixInverse=numpy.delete(FundamentalMatrixInverse,casesToRemove,0)
#FundamentalMatrixInverse=numpy.delete(FundamentalMatrixInverse,casesToRemove,1)
#Boltz=numpy.delete(Boltz,casesToRemove)
#posX=numpy.delete(posX,casesToRemove)
#posY=numpy.delete(posY,casesToRemove)
#posZ=numpy.delete(posZ,casesToRemove)
#protonNumber=numpy.delete(protonNumber,casesToRemove)

# reduce numberOfBindingSites by the number of cases to remove 
#numberOfBindingSites -= len(casesToRemove)

print("number of binding sites")
print(numberOfBindingSites)
print("protons kept ")
print(protonNumber+1)
#print(FundamentalMatrixInverse)

#Fundamental Matrix
FundamentalMatrix=numpy.linalg.inv(FundamentalMatrixInverse)

#print FundamentalMatrix
AvgTime=numpy.zeros(numberOfBindingSites)
R=numpy.zeros((numberOfBindingSites,numberOfBindingSites))
mij=numpy.zeros((numberOfBindingSites,numberOfBindingSites))
Centrality=numpy.zeros(numberOfBindingSites)
#betweeness or flow through centrality

#need mij first
for i in range(numberOfBindingSites):
  for j in range(numberOfBindingSites):
        mij[i][j]=(FundamentalMatrix[j][j]-FundamentalMatrix[i][j])/Boltz[j]*expectedTimeForFirstStep
        for n in range(numberOfBindingSites):
                mij[i][j]=mij[i][j]+(FundamentalMatrix[i][n]-FundamentalMatrix[j][n])*cVector[n]


#next calculate flow-through centrality
for t in range(numberOfBindingSites):
  AvgTime[t]=0.0
  for i in range(numberOfBindingSites):
        for j in range(numberOfBindingSites):
                R[i][j]=mij[i][t]+mij[t][j]
                if mij[i][j]>0:
                        AvgTime[t]=AvgTime[t]+R[i][j]/mij[i][j]
                #print(mij[i][j], mij[i][t], mij[t][j],R[i][j], AvgTime[t])
  Centrality[t]=1/AvgTime[t]
  #print(Centrality[t])

print("AvgTime")
print(AvgTime)
#stretching centrality to cover largest gray scale or largest 0 to 1 range for best color contrast
maxCentrality=numpy.amax(Centrality)
minCentrality=numpy.amin(Centrality)
rangeCentrality=maxCentrality-minCentrality
print("min, max, range of centrality")
print(minCentrality,maxCentrality,rangeCentrality)
print("Centrality")
print(Centrality)
CentralitySpreadOut=(Centrality-minCentrality)/rangeCentrality
print("CentralitySpreadOut")
print(CentralitySpreadOut)

#Reading backbone and proton sites
lattice=numpy.zeros((3,3))
lengths=numpy.zeros(3)
angles=numpy.zeros(3)

numberOfAtomsInContcar=8+7+24+1;#Ba 8 Zr 7 O 24 Y 1 H 1 
contcarCoord=numpy.zeros((numberOfAtomsInContcar+1,3))
numberOfAtomTypes=5
natoms=numpy.zeros(numberOfAtomTypes,dtype=int)
j=0
#need to change file depending on which set you are making
for line in open("POSCAR"):
	if j==1:
		dummy=line.split()
		latFactor=eval(dummy[0])
		print(latFactor)
	if j>1 and j <5:
		dummy=line.split()
		lattice[j-2][0]=latFactor*eval(dummy[0])
		lattice[j-2][1]=latFactor*eval(dummy[1])
		lattice[j-2][2]=latFactor*eval(dummy[2])
		lengths[j-2]=numpy.linalg.norm(lattice[j-2])
	if j==5:
		print(lengths)
		print(lattice)
		angles[0]=numpy.arccos(numpy.dot(lattice[0],lattice[1])/lengths[0]/lengths[1])*180/numpy.pi
		angles[1]=numpy.arccos(numpy.dot(lattice[1],lattice[2])/lengths[2]/lengths[1])*180/numpy.pi
		angles[2]=numpy.arccos(numpy.dot(lattice[0],lattice[2])/lengths[0]/lengths[2])*180/numpy.pi
		print(angles)
		dummy=line.split()
	if j==6:
		dummy=line.split()
		for nn in range(0,numberOfAtomTypes):
		  natoms[nn]=eval(dummy[nn])	
	if (j>7 and j<=7+numberOfAtomsInContcar):
		dummy=line.split()
		contcarCoord[j-7][0]=eval(dummy[0])
		contcarCoord[j-7][1]=eval(dummy[1])
		contcarCoord[j-7][2]=eval(dummy[2])
		#print(j-7,contcarCoord[j-7][0],contcarCoord[j-7][1],contcarCoord[j-7][2])
	j=j+1

#setting up all protons types in original scenario including site 35
natoms[numberOfAtomTypes-1]=96
print(natoms)

#creating vesta file with objects - binding sites and transition states
#atomLabel=["Ba","Zr","O","Al","H"]
#atomDescription=["  2.2400 190 101 116  30 239  44 204  0","  1.6000   0 255   0   0 255   0 204  0","  0.7400 254   3   0 254   3   0 204  0"," 1.4300 129 178 214 129 178 214 204  0"," 0.4600 "]
#atomLabel=["Ba","Zr","O","Sc","H"]
#atomDescription=["  2.2400 190 101 116  30 239  44 204  0","  1.6000   0 255   0   0 255   0 204  0","  0.7400 254   3   0 254   3   0 204  0"," 1.6400 181 99 171 181 99 171 204  0"," 0.4600 "]
atomLabel=["Ba","Zr","O","Y","H"]
atomDescription=["  2.2400 190 101 116  30 239  44 204  0","  1.6000   0 255   0   0 255   0 204  0","  0.7400 254   3   0 254   3   0 204  0"," 1.8200 102 152 142 102 152 142 204  0"," 0.4600 "]

nameoffile="CentralityFlowThrough.vesta"
vesta=open(nameoffile,'w')
vesta.write("#VESTA_FORMAT_VERSION 3.3.0\nCRYSTAL\nTITLE\nAllProtonsColors\nGROUP\n1 1 P 1\nSYMOP\n")
vesta.write(" 0.000000  0.000000  0.000000  1  0  0   0  1  0   0  0  1   1\n")
vesta.write(" -1.0 -1.0 -1.0  0 0 0  0 0 0  0 0 0\nTRANM 0\n")
vesta.write(" 0.000000  0.000000  0.000000  1  0  0   0  1  0   0  0  1\nLTRANSL\n -1\n")
vesta.write(" 0.000000  0.000000  0.000000  0.000000  0.000000  0.000000\n")
vesta.write("LORIENT\n -1   0   0   0   0\n 1.000000  0.000000  0.000000  1.000000  0.000000  0.000000\n")
vesta.write(" 0.000000  0.000000  1.000000  0.000000  0.000000  1.000000\nLMATRIX\n")
vesta.write(" 1.000000  0.000000  0.000000  0.000000\n 0.000000  1.000000  0.000000  0.000000\n")
vesta.write(" 0.000000  0.000000  1.000000  0.000000\n 0.000000  0.000000  0.000000  1.000000\n")
vesta.write(" 0.000000  0.000000  0.000000\n")
vesta.write("CELLP\n  "+str(lengths[0])+" "+str(lengths[1])+" "+str(lengths[2])+" "+str(angles[1])+" "+str(angles[2])+" "+str(angles[0])+"\n")
vesta.write("  0.000000   0.000000   0.000000   0.000000   0.000000   0.000000\nSTRUC\n")
atomStart=1
iproton=0
for itype in range(0,numberOfAtomTypes):
  atomEnd=atomStart+natoms[itype]
  for nn in range(atomStart,atomEnd):
    if itype==numberOfAtomTypes-1:
        vesta.write("  "+str(atomStart+iproton)+"  "+atomLabel[itype]+" "+atomLabel[itype]+str(nn-atomStart+1)+"  1.0000   "+str(posX[iproton])+" "+str(posY[iproton])+" "+str(posZ[iproton])+"    1a       1\n")
        vesta.write("                            0.000000   0.000000   0.000000  0.00\n")
        iproton=iproton+1
    else:
      vesta.write("  "+str(nn)+"  "+atomLabel[itype]+" "+atomLabel[itype]+str(nn-atomStart+1)+"  1.0000   "+str(contcarCoord[nn][0])+" "+str(contcarCoord[nn][1])+" "+str(contcarCoord[nn][2])+"    1a       1\n")
      vesta.write("                            0.000000   0.000000   0.000000  0.00\n")
  atomStart=atomEnd
#after atoms
vesta.write("                            0.000000   0.000000   0.000000  0.00\n")
vesta.write("  0 0 0 0 0 0 0\nTHERI 0\n")
atomStart=1
iproton=0
for itype in range(0,numberOfAtomTypes):
  atomEnd=atomStart+natoms[itype]
  for nn in range(atomStart,atomEnd):
    if itype==numberOfAtomTypes-1:
        vesta.write(" "+str(atomStart+iproton)+"  "+atomLabel[itype]+str(nn-atomStart+1)+"  1.000000\n")
        iproton=iproton+1
    else:
    	vesta.write(" "+str(nn)+"  "+atomLabel[itype]+str(nn-atomStart+1)+"  1.000000\n")
  atomStart=atomEnd
vesta.write("  0 0 0\nSHAPE\n  0       0       0       0   0.000000  0   192   192   192   192\n")
vesta.write("BOUND\n    -0.1      1.1      -0.1      1.1      -0.1      1.1\n  0   0   0   0  0\n")
vesta.write("SBOND\n  1     O     H    0.00000    1.50000  0  1  0  0  1  0.250  2.000 127 127 127\n")
vesta.write("  2     Al     O    0.00000    2.10740  0  1  1  0  1  0.250  2.000 127 127 127\n")
vesta.write("  2     Sc     O    0.00000    2.42029  0  1  1  0  1  0.250  2.000 127 127 127\n")
vesta.write("  2     Y     O    0.00000    2.62549  0  1  1  0  1  0.250  2.000 127 127 127\n")
vesta.write("  3    Zr     O    0.00000    3.09651  0  1  1  0  1  0.250  2.000 127 127 127\n")
vesta.write("  0 0 0 0\nSITET\n")
atomStart=1
iproton=0
for itype in range(0,numberOfAtomTypes):
  atomEnd=atomStart+natoms[itype]
  for nn in range(atomStart,atomEnd):
    if itype==numberOfAtomTypes-1:
        greyScaleValue=int(255*(1-CentralitySpreadOut[iproton])) #255 255 255 is white (least central) 0 0 0 is black (most central)
        threeGreyScale=str(greyScaleValue)+" "+str(greyScaleValue)+" "+str(greyScaleValue)
        vesta.write(" "+str(atomStart+iproton)+"  "+atomLabel[itype]+str(nn-atomStart+1)+atomDescription[itype]+threeGreyScale+" "+threeGreyScale+" 204  0\n")
        iproton=iproton+1
    else:
      vesta.write(" "+str(nn)+"  "+atomLabel[itype]+str(nn-atomStart+1)+atomDescription[itype]+"\n")
  #print(atomStart,natoms[itype],atomEnd)
  atomStart=atomEnd
vesta.write("  0 0 0 0 0 0\nVECTR\n 0 0 0 0 0\nVECTT\n 0 0 0 0 0\nSPLAN\n  0   0   0   0\n")
vesta.write("LBLAT\n -1\nLBLSP\n -1\n")
vesta.write("DLATM\n0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63  -1\nDLBND\n -1\nDLPLY\n -1\n")
vesta.write("PLN2D\n  0   0   0   0\nATOMT\n  1         Ba  2.2400 190 101 116  30 239  44 204\n")
vesta.write("  2         Zr  1.6000   0 255   0   0 255   0 204\n  3          O  0.7400 254   3   0 254   3   0 204\n")
#vesta.write("  4          Al  1.4300 129 178 214 129 178 214 204\n  5          H  0.4600 255 204 204 255 204 204 204\n  0 0 0 0 0 0\n")
#vesta.write("  4          Sc  1.6400 181 99 171 181 99 171 204\n  5          H  0.4600 255 204 204 255 204 204 204\n  0 0 0 0 0 0\n")
vesta.write("  4          Y  1.8200 102 152 142 102 152 142 204\n  5          H  0.4600 255 204 204 255 204 204 204\n  0 0 0 0 0 0\n")
vesta.write("SCENE\n-0.990443 -0.074351  0.116164  0.000000\n-0.070893  0.996917  0.033633  0.000000\n")
vesta.write("-0.118306  0.025075 -0.992660  0.000000\n 0.000000  0.000000  0.000000  1.000000\n")
vesta.write("  0.000   0.000\n  0.000\n  1.425\nHBOND 0 2\n\n")
vesta.write("STYLE\nDISPF 39850946\nMODEL   0  1  0\nSURFS   0  1  1\nSECTS  96  1\n")
vesta.write("FORMS   0  1\nATOMS   0  0  1\nBONDS   1\nPOLYS   1\nVECTS 1.000000\n")
vesta.write("FORMP\n  1  1.0   0   0   0\nATOMP\n 24  24   0  50  2.0   0\nBONDP\n")
vesta.write("  1  16  0.250  2.000 127 127 127\nPOLYP\n 204 1  1.000 180 180 180\n")
vesta.write("ISURF\n  0   0   0   0\nTEX3P\n  1 0.00000E+000 1.00000E+000\nSECTP\n")
vesta.write("  1 0.00000E+000 1.00000E+000 0.00000E+000\nHKLPP\n 192 1  1.000 255   0 255\n")
vesta.write("UCOLP\n   0   1  1.000   0   0   0\nCOMPS 1\nLABEL 1    12  1.000 0\nPROJT 1  0.962\n")
vesta.write("BKGRC\n 255 255 255\nDPTHQ 1 -0.5000  3.5000\nLIGHT0 1\n")
vesta.write(" 1.000000  0.000000  0.000000  0.000000\n 0.000000  1.000000  0.000000  0.000000\n")
vesta.write(" 0.000000  0.000000  1.000000  0.000000\n 0.000000  0.000000  0.000000  1.000000\n")
vesta.write(" 0.000000  0.000000 20.000000  0.000000\n 0.000000  0.000000 -1.000000\n  26  26  26 255\n")
vesta.write(" 179 179 179 255\n 255 255 255 255\nLIGHT1\n 1.000000  0.000000  0.000000  0.000000\n")
vesta.write(" 0.000000  1.000000  0.000000  0.000000\n 0.000000  0.000000  1.000000  0.000000\n")
vesta.write(" 0.000000  0.000000  0.000000  1.000000\n 0.000000  0.000000 20.000000  0.000000\n")
vesta.write(" 0.000000  0.000000 -1.000000\n   0   0   0   0\n   0   0   0   0\n   0   0   0   0\n")
vesta.write("LIGHT2\n 1.000000  0.000000  0.000000  0.000000\n 0.000000  1.000000  0.000000  0.000000\n")
vesta.write(" 0.000000  0.000000  1.000000  0.000000\n 0.000000  0.000000  0.000000  1.000000\n")
vesta.write(" 0.000000  0.000000 20.000000  0.000000\n 0.000000  0.000000 -1.000000\n")
vesta.write("   0   0   0   0\n   0   0   0   0\n   0   0   0   0\n")
vesta.write("LIGHT3\n 1.000000  0.000000  0.000000  0.000000\n 0.000000  1.000000  0.000000  0.000000\n")
vesta.write(" 0.000000  0.000000  1.000000  0.000000\n 0.000000  0.000000  0.000000  1.000000\n")
vesta.write(" 0.000000  0.000000 20.000000  0.000000\n 0.000000  0.000000 -1.000000\n")
vesta.write("   0   0   0   0\n   0   0   0   0\n   0   0   0   0\n")
vesta.write("ATOMM\n 204 204 204 255\n  25.600\nBONDM\n 255 255 255 255\n 128.000\n")
vesta.write("POLYM\n  255 255 255 255\n 128.000\nSURFM\n   0   0   0 255\n 128.000\n")
vesta.write("FORMM\n 255 255 255 255\n 128.000\nHKLPM\n 255 255 255 255\n 128.000\n")
vesta.close()

